# -*- coding: UTF-8 -*-
import torch
import models

model = models.__dict__['eca_resnet50'](k_size=[3, 5, 5, 7], pretrained=True)
model.eval()
checkpoint = torch.load('eca_resnet50_k3557.pth.tar', map_location='cpu')
model = torch.nn.DataParallel(model)
print('Number of models parameters: {}'.format(
    sum([p.data.nelement() for p in model.parameters()])))
model.load_state_dict(checkpoint['state_dict'])
batch_size = 1
input = torch.randn(batch_size, 3, 224, 224, requires_grad=True)
if isinstance(model, torch.nn.DataParallel):
    model = model.module
torch.onnx.export(model, input, 'ecanet50.onnx', export_params=True,
                  opset_version=10, do_constant_folding=True, input_names=['input'], output_names=['output'],
                  dynamic_axes={'input': {0: '?'}, 'output': {0: '?'}})
